import aiohttp
import base64
from typing import Dict, Any, List

from .ai_base import BaseAIService, GenerationRequest, GenerationResponse, AIServiceProvider, AIServiceType, TaskStatus

class StabilityAIService(BaseAIService):
    """Stability AI (Stable Diffusion) 服务实现"""
    
    def __init__(self, config: Dict[str, Any]):
        super().__init__(config)
        self.api_key = config.get("api_key")
        self.base_url = config.get("base_url", "https://api.stability.ai/v1")
        
    def get_provider_name(self) -> AIServiceProvider:
        return AIServiceProvider.STABILITY_AI
    
    def get_supported_services(self) -> List[AIServiceType]:
        return [
            AIServiceType.IMAGE_GENERATION, 
            AIServiceType.IMAGE_EDITING
        ]
    
    async def generate_image(self, request: GenerationRequest) -> GenerationResponse:
        """调用Stability AI生成图片"""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
            "Accept": "application/json"
        }
        
        # 构建Stability AI请求
        stability_request = {
            "text_prompts": [
                {
                    "text": request.prompt,
                    "weight": 1.0
                }
            ],
            "cfg_scale": request.params.get("cfg_scale", 7),
            "height": request.params.get("height", 1024),
            "width": request.params.get("width", 1024),
            "steps": request.params.get("steps", 50),
            "samples": request.params.get("samples", 1),
            "sampler": request.params.get("sampler", "K_DPM_2_ANCESTRAL"),
        }
        
        # 添加负面提示词
        if "negative_prompt" in request.params:
            stability_request["text_prompts"].append({
                "text": request.params["negative_prompt"],
                "weight": -1.0
            })
        
        engine_id = request.params.get("engine_id", "stable-diffusion-v1-6")
        
        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    f"{self.base_url}/generation/{engine_id}/text-to-image",
                    headers=headers,
                    json=stability_request
                ) as response:
                    if response.status == 200:
                        result = await response.json()
                        
                        # Stability AI返回base64编码的图片
                        image_data = result["artifacts"][0]["base64"]
                        
                        # TODO: 上传到对象存储并获取URL
                        # 这里暂时返回一个模拟URL
                        result_url = f"https://storage.thrush-ai.com/images/stability_{hash(image_data[:100])}.png"
                        
                        return GenerationResponse(
                            success=True,
                            result_url=result_url,
                            metadata={
                                "provider": self.provider,
                                "engine": engine_id,
                                "seed": result["artifacts"][0]["seed"],
                                "cfg_scale": stability_request["cfg_scale"],
                                "steps": stability_request["steps"]
                            },
                            cost=self.calculate_cost(request)
                        )
                    else:
                        error_data = await response.json()
                        return GenerationResponse(
                            success=False,
                            error_message=error_data.get("message", "Stability AI API调用失败")
                        )
                        
        except Exception as e:
            return GenerationResponse(
                success=False,
                error_message=f"Stability AI调用失败: {str(e)}"
            )
    
    async def generate_video(self, request: GenerationRequest) -> GenerationResponse:
        """Stability AI的视频生成服务"""
        # Stability AI的Stable Video Diffusion
        headers = {
            "Authorization": f"Bearer {self.api_key}",
        }
        
        # 需要一个起始图片来生成视频
        if "image_url" not in request.params:
            return GenerationResponse(
                success=False,
                error_message="视频生成需要提供起始图片"
            )
        
        # TODO: 实现Stable Video Diffusion API调用
        return GenerationResponse(
            success=False,
            error_message="Stable Video Diffusion功能待实现"
        )
    
    async def edit_image(self, request: GenerationRequest, image_url: str) -> GenerationResponse:
        """使用Stability AI编辑图片"""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Accept": "application/json"
        }
        
        # Image-to-Image API
        stability_request = {
            "text_prompts": [{"text": request.prompt, "weight": 1.0}],
            "cfg_scale": request.params.get("cfg_scale", 7),
            "steps": request.params.get("steps", 50),
            "image_strength": request.params.get("image_strength", 0.35),
        }
        
        # TODO: 下载原图片并作为multipart上传
        return GenerationResponse(
            success=False,
            error_message="图片编辑功能需要实现文件处理逻辑"
        )
    
    async def face_swap(self, request: GenerationRequest, source_url: str, target_url: str) -> GenerationResponse:
        """Stability AI不直接支持人脸替换"""
        return GenerationResponse(
            success=False,
            error_message="Stability AI不支持人脸替换功能，建议使用专门的人脸替换服务"
        )
    
    async def get_task_status(self, task_id: str) -> TaskStatus:
        """Stability AI是同步API"""
        return TaskStatus.COMPLETED
    
    def calculate_cost(self, request: GenerationRequest) -> float:
        """计算Stability AI成本"""
        # 基于steps计算成本
        steps = request.params.get("steps", 50)
        base_cost = 0.01  # 基础成本
        return base_cost * (steps / 50)  # steps越多成本越高